#from Dialogs import *
from Parameters import *
from PolymorphismTables import PolymorphismTables
from R2_Tools import R2
from SiteFrequencySpectrum import SiteFrequencySpectrum
import infiniteSitesTools as ist 
from Polyspect_PowerSphere import PolyspectPowerSphere
from Databases import DB_Polyspect

import math
#import matplotlib.pyplot as plt
import numpy.matlib as np

def main():
	
	print "***************************************************"
	print "********************* New run *********************"
	print "***************************************************"
	print ""

	showDialogs = False
	
	ptNull = None
	mySFS = None

	alphaToCheck = [1.0] #, 2.0, 5.0, 10.0, 20.0, 50.0, 100.0]
	tauToCheck = [10000, 5000, 2000, 1000, 500, 200, 100]

	for alpha in alphaToCheck:
		 for tau in tauToCheck:
			if showDialogs:
				diaPop = Dialog_Population()
				pop = diaPop.getResponse()
			
			else:
				pop = getDefaultPopulationParameters()
			
			nullDemography = Demography( pop.N )
			nullDemography.addEpoch( pop.N, 0, True )
			print nullDemography
			print pop.getMSString_Null(), nullDemography.getMSString()
			
			if showDialogs:
				modelPop = Dialog_MEMES()
				memes = modelPop.getResponse()
				
				demDialog = None
				
				if memes == modelPop.memes[0]:
					demDialog = Dialog_Demography_C( pop.N )
				elif memes == modelPop.memes[1]:
					demDialog = Dialog_Demography_G( pop.N )
				elif memes == modelPop.memes[2]:
					demDialog = Dialog_Demography_CC( pop.N )
				elif memes == modelPop.memes[3]:
					demDialog = Dialog_Demography_CG( pop.N )
				else:
					return
					
				altDemography =  demDialog.getResponse()
			
			else:
				N1 = pop.N * math.exp(-alpha * tau / (4.0*pop.N))
				
				altDemography = Demography(pop.N)
				altDemography.addEpoch( alpha, 0, False)
				altDemography.addEpoch( N1, tau, True )
				
						
			simParams = SimulationParameters( pop, nullDemography )
			print simParams
		
			if ptNull is None:
				ptNull = PolymorphismTables()
				msCommand = simParams.getMSString_Null()
				msOut = ptNull.runMS(msCommand)
				ptNull.readMS(None, msOut)
		
				mySFSs = SiteFrequencySpectrum(ptNull)
			
			pt = PolymorphismTables()
			msCommand = simParams.getMSString_Alternative(altDemography )
			msOut = pt.runMS(msCommand)
			pt.readMS(None, msOut)
			myAltSFSs = SiteFrequencySpectrum(pt)
			
			db = DB_Polyspect()
			polyspects = db.getAllPolyspects( simParams.getNumChromosomes() )
			
			mySpace = PolyspectPowerSphere( simParams, altDemography, mySFSs, myAltSFSs, polyspects)
			
			point = mySpace.getCurrentPoint()
			print point
			
			for i in range(0,10000):
				newPoint = mySpace.getNextPoint()
				if newPoint == None:
					print "\nFound peak for now"
					break
					
				print newPoint
		
			db.addPowerPoint( mySpace )
					
#	rr = R2( simParams )
#	r2power = rr.getPower(altDemography)
#	print 'R2 power:\t' + str(r2power)
#
#	makePlot( mySpace, r2power )


def getDefaultPopulationParameters():
	
	popParameters = Population()
	
	popParameters.numChromosomes = 30
	popParameters.numSegregatingSites = 40

	popParameters.rateMutation = 1E-8
	popParameters.rateRecombination = 0.0
	popParameters.sequenceLength = 25000

	popParameters.N = 10000
	popParameters.numSimulations = 50000
	popParameters.numPowerSimulations = 2500

	return popParameters


def makePlot( mySpace, r2power ):
	tjd = ist.normalizeToFirst( mySpace.myPowerPath[0].myPolyspect.myOperator )
	flf = ist.normalizeToFirst( mySpace.myPowerPath[3].myPolyspect.myOperator )
	dmd = ist.normalizeToFirst( mySpace.myPowerPath[-1].myPolyspect.myOperator )


	myPowers = []
	myPowers.append( mySpace.myPowerPath[0].myPower)
	myPowers.append( mySpace.myPowerPath[2].myPower)
	myPowers.append( mySpace.myPowerPath[3].myPower)
	myPowers.append( r2power )
	myPowers.append( mySpace.myPowerPath[-1].myPower)
	
	fig = plt.figure()
	ind = np.arange(30)
	
	ax = fig.add_subplot(2,2,1)
	ax.bar( ind, tjd, color='m' )
	ax.set_title('Tajimas D')

	ax = fig.add_subplot(2,2,2)
	ax.bar( ind, flf, color='b')
	ax.set_title('Fu Lis F')

	ax = fig.add_subplot(2,2,3)
	ax.bar( ind, dmd, color='y')
	ax.set_title('Davids M')

	ind = np.arange(5)
	ax = fig.add_subplot(2,2,4)
	rects = ax.bar( ind, myPowers, color='r')
	ax.set_title('Powers')
	ax.set_yticks([0, 0.2, 0.4, 0.6, 0.8, 1.0])
	ax.set_xticks(ind+0.35)
	ax.set_xticklabels( ('TajD', 'FuLiD', 'FuLiF', 'R2', 'DavidsM') )

	def autolabel(rects):
		# attach some text labels
		for rect in rects:
			height = rect.get_height()
			ax.text(rect.get_x()+rect.get_width()/2., 1.05*height, '%d'%int(100*height), ha='center', va='bottom')

	autolabel(rects)

	print "\n*** Final path ***"
	print mySpace
	print "\nProgram complete"
	
	plt.show()


	#r = codebox( "*** Final path ***", "Results", str(mySpace) )
	
main()
